"""
MTL Configuration Management
多任务学习配置管理
"""

import json
from typing import Dict, List, Any, Optional, Union
from dataclasses import dataclass, field


@dataclass
class TaskConfig:
    """单个任务的配置"""
    name: str
    task_type: str  # 'classification', 'regression', 'segmentation'
    num_classes: Optional[int] = None  # 分类任务的类别数
    loss_fn: str = 'mse'  # 'mse', 'cross_entropy', 'bce'
    metrics: List[str] = field(default_factory=lambda: ['accuracy'])
    weight: float = 1.0  # 任务权重


@dataclass
class ArchitectureConfig:
    """架构配置"""
    name: str  # 'HPS', 'CrossStitch', 'MTAN', 'MMoE'
    shared_layers: List[int] = field(default_factory=lambda: [512, 256])
    task_specific_layers: List[int] = field(default_factory=lambda: [128, 64])
    
    # 特定架构的参数
    # CrossStitch
    alpha: float = 0.5
    
    # MTAN
    attention_dim: int = 256
    
    # MMoE
    num_experts: int = 4
    expert_dim: int = 256


@dataclass
class WeightingConfig:
    """权重策略配置"""
    name: str  # 'EW', 'GradNorm', 'UW', 'DWA', 'PCGrad', 'CAGrad'
    
    # GradNorm参数
    alpha: float = 0.12
    
    # DWA参数
    temperature: float = 2.0
    
    # PCGrad参数
    reduction: str = 'mean'  # 'mean', 'sum'
    
    # CAGrad参数
    c: float = 0.4
    calpha: float = 0.5


@dataclass
class TrainingConfig:
    """训练配置"""
    batch_size: int = 32
    num_epochs: int = 100
    learning_rate: float = 0.001
    optimizer: str = 'adam'  # 'sgd', 'adam'
    scheduler: Optional[str] = None  # 'step', 'cosine'
    
    # 优化器参数
    momentum: float = 0.9
    weight_decay: float = 1e-4
    
    # 调度器参数
    step_size: int = 30
    gamma: float = 0.1
    
    # 早停
    early_stopping: bool = False
    patience: int = 10
    
    # 梯度裁剪
    grad_clip: Optional[float] = None


class MTLConfig:
    """多任务学习总配置类"""
    
    def __init__(self):
        self.tasks: Dict[str, TaskConfig] = {}
        self.architecture: ArchitectureConfig = ArchitectureConfig(name='HPS')
        self.weighting: WeightingConfig = WeightingConfig(name='EW')
        self.training: TrainingConfig = TrainingConfig()
        self.input_dim: int = -1  # 输入维度
        
    def add_task(self, task_config: TaskConfig):
        """添加任务配置"""
        self.tasks[task_config.name] = task_config
        
    def get_task_names(self) -> List[str]:
        """获取所有任务名称"""
        return list(self.tasks.keys())
    
    def get_num_tasks(self) -> int:
        """获取任务数量"""
        return len(self.tasks)
    
    def set_architecture(self, arch_config: ArchitectureConfig):
        """设置架构配置"""
        self.architecture = arch_config
        
    def set_weighting(self, weight_config: WeightingConfig):
        """设置权重策略配置"""
        self.weighting = weight_config
        
    def set_training(self, train_config: TrainingConfig):
        """设置训练配置"""
        self.training = train_config
        
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'tasks': {name: task.__dict__ for name, task in self.tasks.items()},
            'architecture': self.architecture.__dict__,
            'weighting': self.weighting.__dict__,
            'training': self.training.__dict__,
            'input_dim': self.input_dim,
            'device': self.device
        }
    
    def save(self, filepath: str):
        """保存配置到文件"""
        with open(filepath, 'w') as f:
            json.dump(self.to_dict(), f, indent=2)
    
    @classmethod
    def load(cls, filepath: str) -> 'MTLConfig':
        """从文件加载配置"""
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        config = cls()
        
        # 加载任务配置
        for task_name, task_data in data['tasks'].items():
            task_config = TaskConfig(**task_data)
            config.add_task(task_config)
        
        # 加载架构配置
        config.architecture = ArchitectureConfig(**data['architecture'])
        
        # 加载权重配置
        config.weighting = WeightingConfig(**data['weighting'])
        
        # 加载训练配置
        config.training = TrainingConfig(**data['training'])
        
        config.input_dim = data.get('input_dim', 784)
        config.device = data.get('device', 'cpu')
        
        return config
    
    def validate(self) -> bool:
        """验证配置的有效性"""
        if not self.tasks:
            raise ValueError("至少需要定义一个任务")
        
        # 验证任务配置
        for task_name, task in self.tasks.items():
            if task.task_type == 'classification' and task.num_classes is None:
                raise ValueError(f"分类任务 {task_name} 必须指定 num_classes")
        
        # 验证架构配置
        valid_archs = ['HPS', 'CrossStitch', 'MTAN', 'MMoE']
        if self.architecture.name not in valid_archs:
            raise ValueError(f"不支持的架构: {self.architecture.name}")
        
        # 验证权重策略配置
        valid_weights = ['EW', 'GradNorm', 'UW', 'DWA', 'PCGrad', 'CAGrad']
        if self.weighting.name not in valid_weights:
            raise ValueError(f"不支持的权重策略: {self.weighting.name}")
        
        return True


def create_default_config(task_configs: List[TaskConfig]) -> MTLConfig:
    """创建默认配置"""
    config = MTLConfig()
    
    # 添加任务
    for task_config in task_configs:
        config.add_task(task_config)
    
    # 默认架构：硬参数共享
    config.set_architecture(ArchitectureConfig(name='HPS'))
    
    # 默认权重策略：等权重
    config.set_weighting(WeightingConfig(name='EW'))
    
    # 默认训练配置
    config.set_training(TrainingConfig())
    
    return config


def create_classification_task(name: str, num_classes: int, weight: float = 1.0) -> TaskConfig:
    """创建分类任务配置"""
    return TaskConfig(
        name=name,
        task_type='classification',
        num_classes=num_classes,
        loss_fn='cross_entropy',
        metrics=['accuracy'],
        weight=weight
    )


def create_regression_task(name: str, weight: float = 1.0) -> TaskConfig:
    """创建回归任务配置"""
    return TaskConfig(
        name=name,
        task_type='regression',
        loss_fn='mse',
        metrics=['mse', 'mae'],
        weight=weight
    )


def create_segmentation_task(name: str, num_classes: int, weight: float = 1.0) -> TaskConfig:
    """创建分割任务配置"""
    return TaskConfig(
        name=name,
        task_type='segmentation',
        num_classes=num_classes,
        loss_fn='cross_entropy',
        metrics=['iou', 'pixel_accuracy'],
        weight=weight
    ) 